package com.zw.langchainai.repository;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import jakarta.annotation.Resource;
import org.apache.commons.lang3.StringUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Repository;

import java.time.Duration;
import java.util.Collections;
import java.util.List;

/**
 * 使用redis实现会话记忆持久化
 */
@Repository
public class RedisChatMemoryStore implements ChatMemoryStore {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    /**
     * Retrieves messages for a specified chat memory.
     *
     * @param memoryId The ID of the chat memory.
     * @return List of messages for the specified chat memory. Must not be null. Can be deserialized from JSON using {@link ChatMessageDeserializer}.
     */
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        // 获取会话消息
        String string = stringRedisTemplate.opsForValue().get(memoryId.toString());
        // 转换为List<ChatMessage>
        if (StringUtils.isNoneBlank(string)) {
            return ChatMessageDeserializer.messagesFromJson(string);
        } else {
            return Collections.emptyList();
        }
    }

    /**
     * Updates messages for a specified chat memory.
     *
     * @param memoryId The ID of the chat memory.
     * @param messages List of messages for the specified chat memory, that represent the current state of the {@link ChatMemory}.
     *                 Can be serialized to JSON using {@link ChatMessageSerializer}.
     */
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        // 更新或新增会话消息
        // 1. list转json字符串
        String json = ChatMessageSerializer.messagesToJson(messages);
        // 2. 保存会话消息
        stringRedisTemplate.opsForValue().set(memoryId.toString(), json, Duration.ofDays(1));
    }

    /**
     * Deletes all messages for a specified chat memory.
     *
     * @param memoryId The ID of the chat memory.
     */
    @Override
    public void deleteMessages(Object memoryId) {
        stringRedisTemplate.delete(memoryId.toString());
    }

    public void addMessage(Object memoryId, ChatMessage message) {
        // 添加会话消息
        // 1. 判断是否存在
        if (stringRedisTemplate.hasKey(memoryId.toString())) {
            // 2. 存在,则更新
            List<ChatMessage> messageList = getMessages(memoryId);
            messageList.add(message);
            updateMessages(memoryId, getMessages(memoryId));
        } else {
            // 3. 不存在,则创建
            String messagesToJson = ChatMessageSerializer.messagesToJson(Collections.singletonList(message));
            stringRedisTemplate.opsForValue().set(memoryId.toString(), messagesToJson, Duration.ofDays(1));
        }
        stringRedisTemplate.opsForValue().append(memoryId.toString(), message.toString());
    }
}
